pub mod error;

use self::error::{CallError, Error};
use crate::{
    courier::{
        codec::{ChannelEndCodec, FromSliceDeserializer, ToWriterSerializer},
        RawChannelEnd,
    },
    ClientLocalRequestId, RpcCallRequest, RpcCancelRequest, RpcRequest, RpcResponse,
};
use fnv::FnvHashMap;
use futures::{ready, Sink, Stream};
use pin_project::pin_project;
use serde::{de::DeserializeOwned, Serialize};
use std::{
    future::Future,
    pin::Pin,
    sync::{
        atomic::{AtomicU64, Ordering},
        Arc,
    },
    task::{Context, Poll},
    time::Duration,
};
use tokio::{
    io::{AsyncRead, AsyncWrite},
    sync::{mpsc, oneshot},
};

const RESP_ONESHOT_RECV_LIMIT: usize = 100;
const REQ_ONESHOT_SEND_LIMIT: usize = 100;

enum WrappedRequest<Req, Resp> {
    Call {
        req: RpcCallRequest<Req>,
        resp_chan_tx: oneshot::Sender<RpcResponse<Resp>>,
    },
    Cancel(RpcCancelRequest),
}

impl<Req, Resp> WrappedRequest<Req, Resp> {
    fn extract_inner_req_type(self) -> Option<Req> {
        match self {
            WrappedRequest::Call { req, .. } => Some(req.req),
            WrappedRequest::Cancel(_) => None,
        }
    }
}

#[pin_project]
pub struct RequestingEnd<C, SD, Req, Resp> {
    #[pin]
    raw_end: RawChannelEnd<C, SD, RpcRequest<Req>, RpcResponse<Resp>>,
    send_chan_rx: mpsc::UnboundedReceiver<WrappedRequest<Req, Resp>>,
    awaiting_reqs: FnvHashMap<ClientLocalRequestId, oneshot::Sender<RpcResponse<Resp>>>,

    send_chan_tx: mpsc::UnboundedSender<WrappedRequest<Req, Resp>>,
}

pub fn new_client<C, SD, Req, Resp>(
    c: C,
    serder: SD,
) -> (RequestingEnd<C, SD, Req, Resp>, RequestHandle<Req, Resp>)
where
    C: AsyncRead + AsyncWrite,
    SD: FromSliceDeserializer + ToWriterSerializer,
    Req: Serialize,
    Resp: DeserializeOwned,
{
    let next_id = Arc::new(AtomicU64::new(0));
    let (send_chan_tx, send_chan_rx) = mpsc::unbounded_channel();
    let request_end = RequestingEnd {
        raw_end: RawChannelEnd::new(c, ChannelEndCodec::new(serder)),
        awaiting_reqs: Default::default(),
        send_chan_rx,
        send_chan_tx: send_chan_tx.clone(),
    };
    let handle = RequestHandle {
        next_id,
        send_chan_tx,
    };
    (request_end, handle)
}

enum PollReceiveDone {
    UnderlyingChannelClosed,
    OneshotReceiveLimitReached,
}

enum PollSendDone {
    OneshotSendLimitReached,
}

impl<C, SD, Req, Resp> RequestingEnd<C, SD, Req, Resp>
where
    C: AsyncRead + AsyncWrite,
    SD: FromSliceDeserializer + ToWriterSerializer,
    Req: Serialize,
    Resp: DeserializeOwned,
{
    // Try to send:
    // 1. requests from the request queue
    fn poll_send(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<
        Result<
            PollSendDone,
            Error<<SD as ToWriterSerializer>::Err, <SD as FromSliceDeserializer>::Err>,
        >,
    > {
        for _ in 0..REQ_ONESHOT_SEND_LIMIT {
            ready!(self.as_mut().project().raw_end.poll_ready(cx))?;
            match self.as_mut().project().send_chan_rx.poll_recv(cx) {
                Poll::Ready(Some(req)) => {
                    let req = match req {
                        WrappedRequest::Call { req, resp_chan_tx } => {
                            if resp_chan_tx.is_closed() {
                                continue;
                            }
                            self.as_mut()
                                .project()
                                .awaiting_reqs
                                .insert(req.id, resp_chan_tx);
                            RpcRequest::Call(req)
                        }
                        WrappedRequest::Cancel(req) => {
                            if self
                                .as_mut()
                                .project()
                                .awaiting_reqs
                                .remove(&req.id)
                                .is_none()
                            {
                                continue;
                            }
                            RpcRequest::Cancel(req)
                        }
                    };
                    self.as_mut().project().raw_end.start_send(req)?;
                }
                Poll::Ready(None) => unreachable!("Unexpected state: client send chan closed."),
                Poll::Pending => {
                    ready!(self.as_mut().project().raw_end.poll_flush(cx))?;
                    return Poll::Pending;
                }
            }
        }
        Poll::Ready(Ok(PollSendDone::OneshotSendLimitReached))
    }

    // Try to recv:
    // 1. responses from the underlying channel
    fn poll_recv(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<
        Result<
            PollReceiveDone,
            Error<<SD as ToWriterSerializer>::Err, <SD as FromSliceDeserializer>::Err>,
        >,
    > {
        for _ in 0..RESP_ONESHOT_RECV_LIMIT {
            match ready!(self.as_mut().project().raw_end.poll_next(cx)) {
                Some(Ok(resp)) => {
                    if let Some(tx) = self.as_mut().project().awaiting_reqs.remove(&resp.id()) {
                        // if the recv end is closed, it's fine to drop the response
                        let _ = tx.send(resp);
                    }
                }
                Some(Err(e)) => return Poll::Ready(Err(e.into())),
                None => return Poll::Ready(Ok(PollReceiveDone::UnderlyingChannelClosed)),
            }
        }
        Poll::Ready(Ok(PollReceiveDone::OneshotReceiveLimitReached))
    }
}

impl<C, SD, Req, Resp> Future for RequestingEnd<C, SD, Req, Resp>
where
    C: AsyncRead + AsyncWrite,
    SD: FromSliceDeserializer + ToWriterSerializer,
    Req: Serialize,
    Resp: DeserializeOwned,
{
    type Output =
        Result<(), Error<<SD as ToWriterSerializer>::Err, <SD as FromSliceDeserializer>::Err>>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        loop {
            match (self.as_mut().poll_send(cx)?, self.as_mut().poll_recv(cx)?) {
                (poll_send, Poll::Ready(PollReceiveDone::UnderlyingChannelClosed)) => {
                    if self.awaiting_reqs.is_empty() {
                        break Poll::Ready(Ok(()));
                    }
                    let _ = ready!(poll_send);
                }
                (Poll::Pending, Poll::Pending) => break Poll::Pending,
                _ => {}
            }
        }
    }
}

pub struct RequestHandle<Req, Resp> {
    next_id: Arc<AtomicU64>,
    send_chan_tx: mpsc::UnboundedSender<WrappedRequest<Req, Resp>>,
}

impl<Req, Resp> Clone for RequestHandle<Req, Resp> {
    fn clone(&self) -> Self {
        Self {
            next_id: Arc::clone(&self.next_id),
            send_chan_tx: self.send_chan_tx.clone(),
        }
    }
}

impl<Req, Resp> RequestHandle<Req, Resp> {
    pub async fn call(&self, req: Req, timeout: Option<Duration>) -> Result<Resp, CallError<Req>> {
        struct DoCancelGuard<'a, Req, Resp> {
            do_cancel: bool,
            req_id: ClientLocalRequestId,
            send_chan_tx: &'a mpsc::UnboundedSender<WrappedRequest<Req, Resp>>,
            rx: &'a mut oneshot::Receiver<RpcResponse<Resp>>,
        }
        impl<'a, Req, Resp> Drop for DoCancelGuard<'a, Req, Resp> {
            fn drop(&mut self) {
                if self.do_cancel {
                    self.rx.close();
                    let id = self.req_id;
                    let _ = self
                        .send_chan_tx
                        .send(WrappedRequest::Cancel(RpcCancelRequest { id }));
                }
            }
        }

        let (tx, mut rx) = oneshot::channel::<RpcResponse<Resp>>();
        let id = self.next_id.fetch_add(1, Ordering::Relaxed) as ClientLocalRequestId;

        let req = WrappedRequest::Call {
            req: RpcCallRequest { id, req },
            resp_chan_tx: tx,
        };

        self.send_chan_tx.send(req).map_err(|e| {
            CallError::LocalRequestServingTaskEnded(e.0.extract_inner_req_type().unwrap())
        })?;

        let mut guard = DoCancelGuard {
            do_cancel: true,
            req_id: id,
            send_chan_tx: &self.send_chan_tx,
            rx: &mut rx,
        };

        let resp = if let Some(time) = timeout {
            match tokio::time::timeout(time, &mut guard.rx).await {
                Ok(r) => r,
                Err(_) => return Err(CallError::TimeOut),
            }
        } else {
            (&mut guard.rx).await
        };

        guard.do_cancel = false;

        match resp {
            Ok(RpcResponse::Done { resp, .. }) => Ok(resp),
            Ok(RpcResponse::Errored { .. }) => Err(CallError::Panic),
            Err(_) => Err(CallError::RespLost),
        }
    }
}
